import math
from collections import Counter

def euclidean_distance(x1, x2):  # 计算欧氏距离
    if len(x1) != len(x2):
        raise ValueError("两个样本的特征维度必须一致")
    return math.sqrt(sum([(a - b) ** 2 for a, b in zip(x1, x2)]))

def knn_predict(x, train_data, train_labels, k):
    # 检查输入有效性
    if len(train_data) != len(train_labels):
        raise ValueError("训练数据和标签数量必须一致")
    if k <= 0 or k > len(train_data):
        raise ValueError("k值必须为正且不大于训练样本数量")

    # 计算待预测样本与所有训练样本的距离
    distances = []
    for i in range(len(train_data)):
        dist = euclidean_distance(x, train_data[i])
        distances.append((dist, train_labels[i]))

    # 按距离升序排序，取前k个近邻
    distances.sort(key=lambda x: x[0])
    k_nearest_labels = [label for (dist, label) in distances[:k]]

    # 多数投票决定预测结果
    most_common = Counter(k_nearest_labels).most_common(1)
    return most_common[0][0]

def knn_classify(test_data, train_data, train_labels, k):
    predictions = []
    for x in test_data:
        pred = knn_predict(x, train_data, train_labels, k)
        predictions.append(pred)
    return predictions

if __name__ == "__main__":
    # 训练数据
    train_data = [
        [1.2, 3.1], [1.9, 2.8], [2.3, 3.5],  # 类别0
        [5.4, 6.2], [6.1, 5.8], [5.9, 6.5],  # 类别1
        [8.3, 2.1], [7.9, 1.8], [8.5, 2.5]   # 类别2
    ]
    # 训练标签
    train_labels = [0, 0, 0, 1, 1, 1, 2, 2, 2]

    # 测试数据
    test_data = [
        [2.0, 3.0],  # 预期类别0
        [5.8, 6.0],  # 预期类别1
        [8.0, 2.0]   # 预期类别2
    ]

    # k=3时的预测结果
    k = 3
    predictions = knn_classify(test_data, train_data, train_labels, k)

    # 打印结果
    print(f"k={k}时的预测结果：")
    for i, (test_sample, pred) in enumerate(zip(test_data, predictions)):
        print(f"样本 {test_sample} 预测类别：{pred}")